#!/usr/bin/env python
# coding: utf-8

# In[ ]:


import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report, accuracy_score
from sklearn.preprocessing import StandardScaler
import matplotlib.pyplot as plt
import tushare as ts

# 初始化pro接口
pro = ts.pro_api('1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c')

# 拉取数据
df = pro.daily(**{
    "ts_code": "600030.sh",
    "trade_date": "",
    "start_date": 20210101,
    "end_date": 20250101,
    "offset": "",
    "limit": ""
}, fields=[
    "ts_code",
    "trade_date",
    "open",
    "high",
    "low",
    "close",
    "pre_close",
    "change",
    "pct_chg",
    "vol",
    "amount"
])
# 数据预处理
df = df.sort_values('trade_date')  # 按日期排序
df['trade_date'] = pd.to_datetime(df['trade_date'], format='%Y%m%d')
df.set_index('trade_date', inplace=True)

# 计算技术指标
def calculate_technical_indicators(data):
    # 移动平均线
    data['MA5'] = data['close'].rolling(window=5).mean()
    data['MA10'] = data['close'].rolling(window=10).mean()
    data['MA20'] = data['close'].rolling(window=20).mean()

    # 相对强弱指数(RSI)
    delta = data['close'].diff()
    gain = delta.where(delta > 0, 0)
    loss = -delta.where(delta < 0, 0)

    avg_gain = gain.rolling(window=14).mean()
    avg_loss = loss.rolling(window=14).mean()

    rs = avg_gain / avg_loss
    data['RSI'] = 100 - (100 / (1 + rs))

    # MACD
    exp12 = data['close'].ewm(span=12, adjust=False).mean()
    exp26 = data['close'].ewm(span=26, adjust=False).mean()
    data['MACD'] = exp12 - exp26
    data['Signal_Line'] = data['MACD'].ewm(span=9, adjust=False).mean()

    # 布林带
    data['Upper_Band'] = data['MA20'] + 2 * data['close'].rolling(window=20).std()
    data['Lower_Band'] = data['MA20'] - 2 * data['close'].rolling(window=20).std()

    # 价格变化率
    data['Price_Change'] = data['close'].pct_change()

    # 成交量变化
    data['Volume_Change'] = data['vol'].pct_change()

    return data


df = calculate_technical_indicators(df)

# 创建目标变量 - 未来5天的收益率
df['Future_5day_Return'] = df['close'].shift(-5) / df['close'] - 1
df['Target'] = (df['Future_5day_Return'] > 0).astype(int)  # 1表示上涨，0表示下跌

# 删除缺失值
df = df.dropna()

# 特征选择
features = ['open', 'high', 'low', 'close', 'vol', 'MA5', 'MA10', 'MA20',
            'RSI', 'MACD', 'Signal_Line', 'Upper_Band', 'Lower_Band',
            'Price_Change', 'Volume_Change']

X = df[features]
y = df['Target']

# 数据标准化
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X_scaled, y, test_size=0.2, shuffle=False)

# 构建随机森林模型
rf_model = RandomForestClassifier(n_estimators=100,
                                  random_state=42,
                                  max_depth=5,
                                  min_samples_split=5,
                                  min_samples_leaf=2)

rf_model.fit(X_train, y_train)

# 模型预测
y_pred = rf_model.predict(X_test)
y_pred_proba = rf_model.predict_proba(X_test)[:, 1]  # 上涨概率

# 模型评估
print("模型评估报告:")
print(classification_report(y_test, y_pred))
print(f"准确率: {accuracy_score(y_test, y_pred):.2f}")

# 计算策略收益率
test_dates = df.index[-len(y_test):]
test_data = df.loc[test_dates].copy()
test_data['Pred_Signal'] = y_pred  # 预测信号

# 策略收益率计算
test_data['Strategy_Return'] = test_data['Future_5day_Return'] * test_data['Pred_Signal']
test_data['Buy_Hold_Return'] = test_data['Future_5day_Return']

# 累计收益率
test_data['Cum_Strategy_Return'] = (1 + test_data['Strategy_Return']).cumprod()
test_data['Cum_Buy_Hold_Return'] = (1 + test_data['Buy_Hold_Return']).cumprod()

# 绘制收益率曲线
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False
plt.figure(figsize=(12, 6))
plt.plot(test_data['Cum_Strategy_Return'], label='策略收益率')
plt.plot(test_data['Cum_Buy_Hold_Return'], label='买入持有收益率')
plt.title('策略表现对比')
plt.xlabel('日期')
plt.ylabel('累计收益率')
plt.legend()
plt.grid()
plt.show()

# 计算关键指标
total_strategy_return = test_data['Cum_Strategy_Return'][-1] - 1
total_buy_hold_return = test_data['Cum_Buy_Hold_Return'][-1] - 1
win_rate = (test_data['Strategy_Return'] > 0).mean()

print("\n策略表现总结:")
print(f"策略总收益率: {total_strategy_return:.2%}")
print(f"买入持有总收益率: {total_buy_hold_return:.2%}")
print(f"策略胜率: {win_rate:.2%}")
print(f"策略信号比例: {test_data['Pred_Signal'].mean():.2%}")

